import os
from PIL import Image
import glob

def optimize_png(input_path, output_path, max_size_mb=5, quality=85):
    """
    优化PNG图片：转换颜色模式为索引色，逐步缩小尺寸直至文件大小符合要求
    :param input_path: 输入文件路径
    :param output_path: 输出文件路径
    :param max_size_mb: 最大允许的文件大小（MB）
    """
    max_size = max_size_mb * 1024 * 1024  # 转换为字节
    scale_factor = 0.9  # 每次缩小比例
    min_dimension = 100  # 最小允许尺寸

    try:
        # 验证输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在：{input_path}")

        with Image.open(input_path) as img:
            # # 处理透明度通道
            # if img.mode == 'RGBA':
            #     # 创建白色背景
            #     background = Image.new('RGB', img.size, (255, 255, 255))
            #     background.paste(img, mask=img.split()[3])
            #     img = background
            # elif img.mode not in ['RGB', 'P']:
            #     img = img.convert('RGB')

            # 转换索引色
            if img.mode != 'P':
                img = img.convert('P', palette=Image.ADAPTIVE, colors=256)

            width, height = img.size
            original_size = os.path.getsize(input_path)
            print(f"原始尺寸：{width}x{height} | 原始大小：{original_size/1024:.1f}KB")

            # 保存参数
            save_args = {
                'optimize': True,
                'compress_level': 9
            }

            while True:
                try:
                    img.save(output_path, **save_args)
                except Exception as save_error:
                    raise RuntimeError(f"保存失败: {str(save_error)}") from save_error

                current_size = os.path.getsize(output_path)
                
                # 检查文件大小
                if current_size <= max_size:
                    print("大小符合要求")
                    break
                
                # 计算新尺寸（确保最小尺寸）
                new_width = max(min_dimension, int(width * scale_factor))
                new_height = max(min_dimension, int(height * scale_factor))
                
                if new_width == width and new_height == height:
                    print("已达到最小尺寸限制")
                    break

                print(f"调整尺寸至 {new_width}x{new_height}")
                img = img.resize((new_width, new_height), Image.LANCZOS)
                width, height = new_width, new_height

    except Exception as e:
        print(f"处理错误: {str(e)}")
        if os.path.exists(output_path):
            os.remove(output_path)
            print("已清理失败文件")

def batch_process(input_dir):
    output_dir=input_dir+"/upload"
    """批量处理目录下所有PNG文件"""
    try:
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建输出目录：{output_dir}")

        png_files = glob.glob(os.path.join(input_dir, "*.png"))
        total = len(png_files)
        
        print(f"发现 {total} 个PNG文件")
        print("开始转换处理...")
        print("-" * 40)

        for i, input_path in enumerate(png_files, 1):
            filename = os.path.basename(input_path)
            output_path = os.path.join(output_dir, filename)
            
            print(f"\n处理中 ({i}/{total}): {filename}")
            optimize_png(input_path, output_path)
            
            if os.path.exists(output_path):
                final_size = os.path.getsize(output_path) / 1024
                print(f"✔ 处理完成 | 最终大小：{final_size:.1f}KB")
            else:
                print("✖ 处理失败")

    except Exception as e:
        print(f"批量处理失败: {str(e)}")

if __name__ == "__main__":
    batch_process("Z:/跨境/男装/0-2023秋季/5-短袖T/IP款图-0228/3.11-大力水手19款/新建文件夹")